BedrockのKnowledge base、OpenSearch Serverless、S3、Lambdaを利用し、RAGを実装してみた #AWSreInvent

こんにちは、つくぼし(tsukuboshi0755)です！

AWS re:Invent 2023で、Bedrockの新機能であるKnowledge baseが発表されました。

この機能を用いる事で、Bedrockを用いたRAGシステムを簡単に構築する事ができるようになります。

今回はこのKnowledge baseと、OpenSearch Serverless、S3及びLambdaを使用し、API経由で使用可能なRetrieval-Augmented Generation(以下RAG)を実装してみたいと思います！

構成

今回実装する構成は以下の通りです。

上記の構成を、2023年12月時点でKnowledge baseが使用可能なバージニアリージョン(us-east-1)で作成します。

なお今回の構成では、Lambdaにプロンプトを直接渡しています。

実際のシステムでは、API Gateway等を用いて、プロンプトを入力するインターフェース(チャットアプリやウェブフォーム等)とLambdaとの連携が必要になるためご注意ください。

構築手順

S3バケットの作成

事前にKnowledge baseのデータソースとして指定するS3バケットを、以下の通り作成しておきます。

Bedrock Knowledge base(+ OpenSearch Serverless)の作成

以下のブログにおけるKnowledge baseの作成画面から確認画面までの手順を参考に、BedrockのKnowledge baseを作成します。

データソースの設定では、先ほど作成したS3バケットを選択してください。

またVector storeの設定では、Quick create a new vector storeを選択し、新規でOpenSearch Serverlessを作成してください。

作成が完了すると、Knowledge baseの詳細画面にアクセスできるようになります。

またOpenSearch Serverlessのコンソールに移動すると、Knowledge base用のコレクションも合わせて作成されている事が分かります。

なおKnowledge baseを削除しても、OpenSearch Serverlessリソースは残るので、検証が完了した後は消し忘れにご注意ください。

Lambdaの作成

Knowledge baseをAPI経由で呼び出すLambdaを作成します。

以下を参考に、IAMロールにBedrockの権限を付与し、Bedrockを利用するためのBoto3ライブラリをアップロードしてください。

今回使用するLambdaのコードは以下の通りです。

import boto3
import json

bedrock_agent_runtime_client = boto3.client('bedrock-agent-runtime')


# Lambda のハンドラー関数
def lambda_handler(event, context):
    user_prompt = event.get('user_prompt')
    
    # Knowledge Base IDに置き換えてください
    knowledge_base_id = '<Knowledge Base ID>'

    # BedrockのモデルにはClaude V2を利用します
    modelArn = 'arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2'

    prompt = f"""\n\nHuman:
    [質問]に適切に答えてください。回答は日本語に翻訳してください。  
    [質問]
    {user_prompt}
    Assistant:
    """

    response = bedrock_agent_runtime_client.retrieve_and_generate(
        input={
            'text': prompt,
        },
        retrieveAndGenerateConfiguration={
            'type': 'KNOWLEDGE_BASE',
            'knowledgeBaseConfiguration': {
                'knowledgeBaseId': knowledge_base_id,
                'modelArn': modelArn,
            }
        }
    )

    print("Received response:" + json.dumps(response, ensure_ascii=False))

    response_output = response['output']['text']

    return response_output

なお現状Knowledge baseは言語を指定する機能がないためか、LambdaからAPIを叩くと英語で回答が返される事があります。

今回日本語で回答を得るために、プロンプトに「回答は日本語に翻訳してください」という一文を新たに追加しています。

RAGの動作確認

構築したRAGシステムが、S3にアップロードされたPDFファイルを参照し、正しい回答を返答できるか確認します。

Bedrockで使用可能なモデルの確認

まず前提として、対象リージョンで、BedrockのClaudeにアクセス可能か確認してください。

もしアクセスしたいモデルのステータスがAccess grantedになっていなければ、別途モデルアクセスの申請を実施しておきましょう。

S3バケットへのデータ投入

次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。

# S3 バケット名を設定
BUCKET_NAME=cm-test-knowledge-base-bucket-714547411825

# AWS の公式ドキュメントの PDF ファイルをダウンロード
mkdir knowledge_base
cd knowledge_base
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/kendra/latest/dg/kendra-dg.pdf -O Kendra.pdf
wget https://docs.aws.amazon.com/ja_jp/Route53/latest/DeveloperGuide/route53-dg.pdf -O Route53.pdf
cd ..

# S3 バケットに PDF ファイルをアップロード
aws s3 cp knowledge_base s3://${BUCKET_NAME} --recursive

なお今回の確認で使用する文書データの例として、以下の記事で利用しているPDFファイルを流用しています。

ファイルのアップロードが成功すると、以下の通り作成したS3バケット配下に5つのPDFファイルが作成されます。

Knowledge baseのデータ同期

次に作成したKnowledge baseのデータソースに移動し、Syncボタンを押すと、S3にアップロードしたPDFファイルがKnowledge baseに同期されます。

データの同期が成功すると、Status項目がReadyになり、Last Sync Time項目に最後に同期した時間が表示されます。

Lambdaのテスト

最後に作成したLambda関数をクリックし、コードタブにてConfigure test eventをクリックし、テストイベントを作成します。

今回イベントJSONには、キーにuser_promptを指定し、値にRAGに対するユーザーからの任意の質問を入力してください。

例えば以下のような形で指定します。

{
  "user_prompt": "Lambda 関数で使用できるメモリの最大値を教えてください。"
}

イベントJSONを変更したら、呼び出すボタンをクリックしてLambda関数を実行しましょう。

実行後の結果で、以下のようにS3にアップロードしたPDFファイルが参照され、想定されるレスポンスが返る事が確認できればOKです！

AWS Lambda のメモリ上限は 10,240 MB です。

最後に

今回は、BedrockのKnowledge baseと、OpenSearch Serverless、S3及びLambdaを使用し、API経由で使用可能なRAGを実装してみました。

Knowledge baseを使用すると、Bedrockから気軽にRAGを利用できるようになるので、ぜひ一度触ってみてはいかがでしょうか。

以上、つくぼし(tsukuboshi0755)でした！